package cn.dataling.ai.filter;

import cn.dataling.ai.common.Response;
import cn.dataling.ai.common.ResponseStatusCodeEnum;
import cn.dataling.ai.component.JacksonComponent;
import cn.dataling.ai.util.LocalCache;
import cn.dataling.ai.component.XcxComponent;
import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.util.ObjectUtils;

import java.io.IOException;

/**
 * @coder taoym
 * @datetime 2021/10/8 16:14
 * @desc 访问过滤器 每分钟访问超过10次 则屏蔽请求
 */
public class RequestLimitFilter implements Filter {


    private final JacksonComponent jacksonComponent;

    private final XcxComponent xcxComponent;

    public RequestLimitFilter(JacksonComponent jacksonComponent, XcxComponent xcxComponent) {
        this.jacksonComponent = jacksonComponent;
        this.xcxComponent = xcxComponent;
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain)
            throws IOException, ServletException {

        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;

        // 检查是否是WebSocket升级请求
        if (isWebSocketUpgradeRequest(request)) {
            String openId = request.getHeader("openId");

            // 验证逻辑
            if (ObjectUtils.isEmpty(openId) || !xcxComponent.isLogin(openId)) {
                response.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
                response.setCharacterEncoding("UTF-8");
                response.setContentType("application/json; charset=utf-8");
                Response<Object> resp = Response.fail(ResponseStatusCodeEnum.OPEN_ID_GET_FAIL, null);
                response.getWriter().write(jacksonComponent.toJson(resp));
                return; // 注意这里不再调用 filterChain.doFilter()
            }

            // 限流逻辑
            Integer times = LocalCache.get("ACCESS:" + openId, Integer.class);
            if (!ObjectUtils.isEmpty(times) && times > 100) {
                response.setCharacterEncoding("UTF-8");
                response.setContentType("application/json; charset=utf-8");
                Response<Object> resp = Response.fail(ResponseStatusCodeEnum.TRY_AGAIN_LATER, null);
                response.getWriter().write(jacksonComponent.toJson(resp));
                return; // 注意这里不再调用 filterChain.doFilter()
            }
        }

        // 非WebSocket请求或验证通过的WebSocket请求继续处理
        filterChain.doFilter(servletRequest, servletResponse);
    }

    private boolean isWebSocketUpgradeRequest(HttpServletRequest request) {
        return "Upgrade".equalsIgnoreCase(request.getHeader("Connection")) &&
                "websocket".equalsIgnoreCase(request.getHeader("Upgrade"));
    }
}
